package service

import (
	"encoding/json"
	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
	"hcy-api/lib/http"
	"hcy-api/lib/jwt"
	"hcy-api/structs/enum"
	"hcy-api/structs/lms"
	"io"
	"lms/global"
	"net/http"
)

var WSService = new(wsService)

type wsService struct {
}

var wsMap = make(map[int64]*websocket.Conn)

func (s wsService) OllamaChat(c *gin.Context) (err error) {
	var upGrader = websocket.Upgrader{
		CheckOrigin: func(r *http.Request) bool {
			return true
		},
	}
	token := c.GetHeader("Sec-WebSocket-Protocol")
	ws, err := upGrader.Upgrade(c.Writer, c.Request, http.Header{
		"Sec-WebSocket-Protocol": []string{token},
	})
	if err != nil {
		global.LmsLog.Errorf("升级为WebSocket协议失败：%v \n", err.Error())
		http.BadWithCode(c, enum.ErrorMap.WsUpgradeErr)
		return
	}
	errChan := make(chan error, 1)
	claims, err := jwt.NewJWT().ParseToken(token)
	if err != nil {
		return err
	}
	wsMap[claims.UserID] = ws
	go func() {
		for {
			var bs []byte
			_, bs, err = ws.ReadMessage()
			if string(bs) == "{\"pageKey\":-1}" {
				continue
			}
			if err != nil {
				delete(wsMap, claims.UserID)
				errChan <- err
				return
			}
			var in lms.StartChat
			err = json.Unmarshal(bs, &in)
			if err != nil {
				delete(wsMap, claims.UserID)
				errChan <- err
				return
			}
			err = KnowledgeService.Chat(claims.UserID, ws, in)
			if err != nil {
				delete(wsMap, claims.UserID)
				errChan <- err
				return
			}
		}
	}()
	err = <-errChan
	if err != io.EOF {
		global.LmsLog.Error(err)
		return err
	}
	return nil
}
